#!/bin/python
# coding=utf-8
"""
Copyright (C) 2018, winewolf@qq.com All rights reserved.
"""

import socket
import hashlib
from Crypto.Cipher import AES, DES


local_server_port = 8003
proxy_server_ip = '127.0.0.1'
proxy_server_port = 8080
long_sock = None
auth_pwd = 'hello baby!'
auth_padding = '89)$Lg'
BUFFER_SIZE = 16 * 1024 * 1024


def md5(_src):
    return hashlib.md5(_src).digest()


def aes_enc(key, plain):
    aes = AES.new(md5(key + auth_padding), AES.MODE_CBC, md5(key))
    pad_len = 16 - len(plain) % 16
    plain += chr(pad_len)*pad_len
    cipher = aes.encrypt(plain)
    return cipher


def aes_dec(key, cipher):
    aes = AES.new(md5(key + auth_padding), AES.MODE_CBC, md5(key))
    plain = aes.decrypt(cipher)
    return plain[:-ord(plain[-1])]


def connect():
    global long_sock
    if long_sock:
        long_sock.close()
        long_sock = None
    long_sock = socket.socket()
    long_sock.connect((proxy_server_ip, proxy_server_port))
    print '[auth] wait random ...'
    auth_random = long_sock.recv(BUFFER_SIZE)
    # auth: .. calculate
    hash_input = auth_random + auth_pwd + auth_padding
    hash_ack = hashlib.md5(hash_input).digest()
    print '[auth] send hash'
    long_sock.send(hash_ack)
    print '[auth] wait result ...'
    auth_ack = long_sock.recv(BUFFER_SIZE)
    auth_result = aes_dec(auth_pwd, auth_ack)
    print '[auth] result = <%r> --> <%s>' % (auth_ack, auth_result)
    if auth_result == 'OK':
        return True
    else:
        long_sock.close()
        long_sock = None
        return False


def run():
    while True:
        from_proxy = long_sock.recv(BUFFER_SIZE)
        print '[proxy] << %r' % from_proxy
        req = aes_dec(auth_pwd, from_proxy)
        print '[server] >> %s' % req

        sock = socket.socket()
        sock.connect(('127.0.0.1', local_server_port))
        sock.send(req)
        ack = sock.recv(BUFFER_SIZE)

        print '[server] << %s' % ack
        to_proxy = aes_enc(auth_pwd, ack)
        print '[server] >> %r' % to_proxy
        long_sock.send(to_proxy)


if connect():
    run()
